import baostock as bs
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression
from joblib import dump



#线性回归方法
def get_stock_data(code, start_date, end_date):
    # 登录证券宝
    lg = bs.login()
    if lg.error_code!= '0':
        print(f"登录失败: {lg.error_code}, {lg.error_msg}")
        exit()
    # 获取股票的历史数据
    rs = bs.query_history_k_data_plus(code,
                                 "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,peTTM,pbMRQ,psTTM,pcfNcfTTM,isST",
                                 start_date=start_date, end_date=end_date,
                                 frequency="d", adjustflag="3")
    if rs.error_code!= '0':
        print(f"获取数据失败，错误码：{rs.error_code}，错误信息：{rs.error_msg}")
        exit()
    # 转为 DataFrame
    data = []
    while (rs.error_code == '0') & rs.next():
        data.append(rs.get_row_data())
    df = pd.DataFrame(data, columns=rs.fields)
    # 登出证券宝
    bs.logout()
    return df


def main():
    # 获取股票数据
    df = get_stock_data("sz.002158", start_date='2010-01-01', end_date='2024-10-31')
    # 数据类型转换
    for col in ['open', 'high', 'low', 'close', 'preclose', 'volume', 'amount', 'pctChg', 'peTTM', 'pbMRQ', 'psTTM', 'pcfNcfTTM', 'turn']:
        df[col] = pd.to_numeric(df[col])
    # 添加下一个交易日的涨幅
    df['next_pctChg'] = df['pctChg'].shift(-1)
    # 处理最后一个 NaN 值
    df.loc[df.index[-1], 'next_pctChg'] = np.nan
    # 提取输入数据和输出数据
    X = df[['open', 'high', 'low', 'close', 'volume', 'amount', 'pctChg']].values
    y = df['next_pctChg'].values
    # 处理 NaN 值
    mask = ~np.isnan(y)
    X = X[mask]
    y = y[mask]
    # 数据标准化
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)
    # 建立线性回归模型
    model = LinearRegression()
    # 模型训练
    model.fit(X_train, y_train)
    # 预测测试集
    y_pred = model.predict(X_test)
    # 保存模型
    dump(model, 'model_line_0119.h5')
    # 预测 2024 年 12 月的数据
    # 这里假设你已经获取了 2024 年 12 月的数据，实际使用时需补充获取数据的代码
    # 示例中假设 2024 年 12 月的数据存储在 X_december 中
    # X_december = get_december_data()
    # X_december_scaled = scaler.transform(X_december)
    # y_pred_december = model.predict(X_december_scaled)
    # 结果展示
    # 这里先使用测试集的数据作为示例展示
    y_pred = model.predict(X_test)
    # 放大预测数据 10 倍
    y_pred = y_pred * 10
    plt.figure(figsize=(10, 6))
    plt.plot(y_test, label='Actual')
    plt.plot(y_pred, label='Predicted (x10)')
    plt.xlabel('Time')
    plt.ylabel('PCT Change')
    plt.title('Actual vs Predicted PCT Change')
    plt.legend()
    plt.show()


if __name__ == "__main__":
    main()